import subprocess
import os
import datetime
import zipfile
from app.settings.config import settings

db = settings.TORTOISE_ORM["connections"]["mysql"]["credentials"]

class MySQLBackup:
    def __init__(self):
        """
        初始化备份类
        """
        self.host = db['host']
        self.user = db['user']
        self.port = db['port']
        self.password = db['password']
        self.database = db['database']
        self.backup_dir = settings.MYSQL_BAK_FOLDER

        # 确保备份目录存在
        if not os.path.exists(self.backup_dir):
            os.makedirs(self.backup_dir)

    def backup(self):
        """
        执行数据库备份并压缩为 ZIP 文件
        :return: 压缩后的备份文件路径
        """
        # 生成备份文件名
        timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
        backup_file = os.path.join(self.backup_dir, f'{self.database}_{timestamp}.sql')
        compressed_file = f'{backup_file}.zip'  # 压缩后的文件名

        try:
            print(f'Backing up database {self.database} to {backup_file}...')
            with open(backup_file, 'w') as f:
                subprocess.run(
                    [
                        'mysqldump',
                        f'--host={self.host}',
                        f'--port={self.port}',
                        f'--user={self.user}',
                        f'--password={self.password}',
                        self.database
                    ],
                    stdout=f,
                    check=True
                )

            # 压缩备份文件
            print(f'Compressing {backup_file} to {compressed_file}...')
            with zipfile.ZipFile(compressed_file, 'w', zipfile.ZIP_DEFLATED) as zipf:
                zipf.write(backup_file, os.path.basename(backup_file))

            # 删除未压缩的 SQL 文件
            os.remove(backup_file)
            print('Backup and compression completed successfully.')
            return compressed_file

        except subprocess.CalledProcessError as e:
            print(f'Error during backup: {e}')
            return None
        except Exception as e:
            print(f'An error occurred: {e}')
            return None

# 示例用法
if __name__ == '__main__':
    # 创建备份对象
    backup = MySQLBackup()
    # 执行备份
    backup_file = backup.backup()
    if backup_file:
        print(f'Compressed backup saved to: {backup_file}')
    else:
        print('Backup failed.')
